{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 20,
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "import torch\n",
    "from copy import deepcopy\n",
    "import numpy as np\n",
    "import xarray as xr\n",
    "import pandas as pd\n",
    "import torch.nn as nn\n",
    "import random\n",
    "from tqdm import tqdm\n",
    "from sklearn.model_selection import train_test_split\n",
    "from torch.utils.data import Dataset, DataLoader\n",
    "from torch.utils.data import TensorDataset, Dataset, ConcatDataset\n",
    "from datasetN import get_dataset_new as get_dataset"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 21,
   "metadata": {},
   "outputs": [],
   "source": [
    "#设置种子\n",
    "def set_seed(seed = 427):\n",
    "    random.seed(seed)\n",
    "    np.random.seed(seed)\n",
    "    os.environ['PYTHONHASHSEED'] = str(seed)\n",
    "    torch.manual_seed(seed)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 24,
   "metadata": {},
   "outputs": [
    {
     "ename": "TypeError",
     "evalue": "from_tuples() got an unexpected keyword argument 'name'",
     "output_type": "error",
     "traceback": [
      "\u001b[0;31m---------------------------------------------------------------------------\u001b[0m",
      "\u001b[0;31mTypeError\u001b[0m                                 Traceback (most recent call last)",
      "\u001b[0;32m<ipython-input-24-05f715571ab3>\u001b[0m in \u001b[0;36m<module>\u001b[0;34m\u001b[0m\n\u001b[0;32m----> 1\u001b[0;31m \u001b[0mcmip\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mget_dataset\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m'cmip'\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m",
      "\u001b[0;32m~/JupyterLab/Tianchi-AIEarth-2021/notebook/datasetN.py\u001b[0m in \u001b[0;36mget_dataset_new\u001b[0;34m(name, debug_mode, small)\u001b[0m\n\u001b[1;32m     97\u001b[0m             ds_list.append(NCDSFactory(\n\u001b[1;32m     98\u001b[0m                 \u001b[0mtrain\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0misel\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0myear\u001b[0m\u001b[0;34m=\u001b[0m\u001b[0mi_list\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m,\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m---> 99\u001b[0;31m                 \u001b[0mlabel\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0misel\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0myear\u001b[0m\u001b[0;34m=\u001b[0m\u001b[0mi_list\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m    100\u001b[0m             ))\n\u001b[1;32m    101\u001b[0m     \u001b[0;32mif\u001b[0m \u001b[0mname\u001b[0m \u001b[0;34m!=\u001b[0m \u001b[0;34m'cmip5'\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n",
      "\u001b[0;32m~/JupyterLab/Tianchi-AIEarth-2021/notebook/datasetN.py\u001b[0m in \u001b[0;36mNCDSFactory\u001b[0;34m(***failed resolving arguments***)\u001b[0m\n\u001b[1;32m     52\u001b[0m     \u001b[0my_end\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mraw_label\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mtail\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0myear\u001b[0m\u001b[0;34m=\u001b[0m\u001b[0;36m1\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mmonth\u001b[0m\u001b[0;34m=\u001b[0m\u001b[0;36m24\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;31m\\\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m     53\u001b[0m                 \u001b[0;34m.\u001b[0m\u001b[0mstack\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m**\u001b[0m\u001b[0;34m{\u001b[0m\u001b[0mindex_col\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m'year'\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0;34m'month'\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m}\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m---> 54\u001b[0;31m     \u001b[0my\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0my\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mmerge\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0my_end\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m     55\u001b[0m     \u001b[0;32mdel\u001b[0m \u001b[0mraw_label\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m     56\u001b[0m     \u001b[0mcollect\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n",
      "\u001b[0;32m~/anaconda3/lib/python3.7/site-packages/xarray/core/dataset.py\u001b[0m in \u001b[0;36mmerge\u001b[0;34m(self, other, overwrite_vars, compat, join, fill_value)\u001b[0m\n\u001b[1;32m   3965\u001b[0m             \u001b[0mcompat\u001b[0m\u001b[0;34m=\u001b[0m\u001b[0mcompat\u001b[0m\u001b[0;34m,\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m   3966\u001b[0m             \u001b[0mjoin\u001b[0m\u001b[0;34m=\u001b[0m\u001b[0mjoin\u001b[0m\u001b[0;34m,\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m-> 3967\u001b[0;31m             \u001b[0mfill_value\u001b[0m\u001b[0;34m=\u001b[0m\u001b[0mfill_value\u001b[0m\u001b[0;34m,\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m   3968\u001b[0m         )\n\u001b[1;32m   3969\u001b[0m         \u001b[0;32mreturn\u001b[0m \u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0m_replace\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m**\u001b[0m\u001b[0mmerge_result\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0m_asdict\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n",
      "\u001b[0;32m~/anaconda3/lib/python3.7/site-packages/xarray/core/merge.py\u001b[0m in \u001b[0;36mdataset_merge_method\u001b[0;34m(dataset, other, overwrite_vars, compat, join, fill_value)\u001b[0m\n\u001b[1;32m    923\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m    924\u001b[0m     return merge_core(\n\u001b[0;32m--> 925\u001b[0;31m         \u001b[0mobjs\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mcompat\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mjoin\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mpriority_arg\u001b[0m\u001b[0;34m=\u001b[0m\u001b[0mpriority_arg\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mfill_value\u001b[0m\u001b[0;34m=\u001b[0m\u001b[0mfill_value\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m    926\u001b[0m     )\n\u001b[1;32m    927\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n",
      "\u001b[0;32m~/anaconda3/lib/python3.7/site-packages/xarray/core/merge.py\u001b[0m in \u001b[0;36mmerge_core\u001b[0;34m(objects, compat, join, combine_attrs, priority_arg, explicit_coords, indexes, fill_value)\u001b[0m\n\u001b[1;32m    609\u001b[0m     \u001b[0mcoerced\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mcoerce_pandas_values\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mobjects\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m    610\u001b[0m     aligned = deep_align(\n\u001b[0;32m--> 611\u001b[0;31m         \u001b[0mcoerced\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mjoin\u001b[0m\u001b[0;34m=\u001b[0m\u001b[0mjoin\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mcopy\u001b[0m\u001b[0;34m=\u001b[0m\u001b[0;32mFalse\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mindexes\u001b[0m\u001b[0;34m=\u001b[0m\u001b[0mindexes\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mfill_value\u001b[0m\u001b[0;34m=\u001b[0m\u001b[0mfill_value\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m    612\u001b[0m     )\n\u001b[1;32m    613\u001b[0m     \u001b[0mcollected\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mcollect_variables_and_indexes\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0maligned\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n",
      "\u001b[0;32m~/anaconda3/lib/python3.7/site-packages/xarray/core/alignment.py\u001b[0m in \u001b[0;36mdeep_align\u001b[0;34m(objects, join, copy, indexes, exclude, raise_on_invalid, fill_value)\u001b[0m\n\u001b[1;32m    428\u001b[0m         \u001b[0mindexes\u001b[0m\u001b[0;34m=\u001b[0m\u001b[0mindexes\u001b[0m\u001b[0;34m,\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m    429\u001b[0m         \u001b[0mexclude\u001b[0m\u001b[0;34m=\u001b[0m\u001b[0mexclude\u001b[0m\u001b[0;34m,\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m--> 430\u001b[0;31m         \u001b[0mfill_value\u001b[0m\u001b[0;34m=\u001b[0m\u001b[0mfill_value\u001b[0m\u001b[0;34m,\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m    431\u001b[0m     )\n\u001b[1;32m    432\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n",
      "\u001b[0;32m~/anaconda3/lib/python3.7/site-packages/xarray/core/alignment.py\u001b[0m in \u001b[0;36malign\u001b[0;34m(join, copy, indexes, exclude, fill_value, *objects)\u001b[0m\n\u001b[1;32m    319\u001b[0m                 \u001b[0;32mif\u001b[0m \u001b[0mjoin\u001b[0m \u001b[0;34m==\u001b[0m \u001b[0;34m\"exact\"\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m    320\u001b[0m                     \u001b[0;32mraise\u001b[0m \u001b[0mValueError\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34mf\"indexes along dimension {dim!r} are not equal\"\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m--> 321\u001b[0;31m                 \u001b[0mindex\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mjoiner\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mmatching_indexes\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m    322\u001b[0m                 \u001b[0;31m# make sure str coords are not cast to object\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m    323\u001b[0m                 \u001b[0mindex\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mmaybe_coerce_to_str\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mindex\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mall_coords\u001b[0m\u001b[0;34m[\u001b[0m\u001b[0mdim\u001b[0m\u001b[0;34m]\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n",
      "\u001b[0;32m~/anaconda3/lib/python3.7/site-packages/pandas/core/indexes/base.py\u001b[0m in \u001b[0;36munion\u001b[0;34m(self, other, sort)\u001b[0m\n\u001b[1;32m   2273\u001b[0m             \u001b[0;32mreturn\u001b[0m \u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0m_union_incompatible_dtypes\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mother\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0msort\u001b[0m\u001b[0;34m=\u001b[0m\u001b[0msort\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m   2274\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m-> 2275\u001b[0;31m         \u001b[0;32mreturn\u001b[0m \u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0m_union\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mother\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0msort\u001b[0m\u001b[0;34m=\u001b[0m\u001b[0msort\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m   2276\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m   2277\u001b[0m     \u001b[0;32mdef\u001b[0m \u001b[0m_union\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mself\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mother\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0msort\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n",
      "\u001b[0;32m~/anaconda3/lib/python3.7/site-packages/pandas/core/indexes/base.py\u001b[0m in \u001b[0;36m_union\u001b[0;34m(self, other, sort)\u001b[0m\n\u001b[1;32m   2347\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m   2348\u001b[0m         \u001b[0;31m# for subclasses\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m-> 2349\u001b[0;31m         \u001b[0;32mreturn\u001b[0m \u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0m_wrap_setop_result\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mother\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mresult\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m   2350\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m   2351\u001b[0m     \u001b[0;32mdef\u001b[0m \u001b[0m_wrap_setop_result\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mself\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mother\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mresult\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n",
      "\u001b[0;32m~/anaconda3/lib/python3.7/site-packages/pandas/core/indexes/base.py\u001b[0m in \u001b[0;36m_wrap_setop_result\u001b[0;34m(self, other, result)\u001b[0m\n\u001b[1;32m   2350\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m   2351\u001b[0m     \u001b[0;32mdef\u001b[0m \u001b[0m_wrap_setop_result\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mself\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mother\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mresult\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m-> 2352\u001b[0;31m         \u001b[0;32mreturn\u001b[0m \u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0m_constructor\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mresult\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mname\u001b[0m\u001b[0;34m=\u001b[0m\u001b[0mget_op_result_name\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mself\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mother\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m   2353\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m   2354\u001b[0m     _index_shared_docs[\n",
      "\u001b[0;31mTypeError\u001b[0m: from_tuples() got an unexpected keyword argument 'name'"
     ]
    }
   ],
   "source": [
    "cmip = get_dataset('cmip')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [
    {
     "ename": "NameError",
     "evalue": "name 'dataset' is not defined",
     "output_type": "error",
     "traceback": [
      "\u001b[0;31m---------------------------------------------------------------------------\u001b[0m",
      "\u001b[0;31mNameError\u001b[0m                                 Traceback (most recent call last)",
      "\u001b[0;32m<ipython-input-1-3075b36c09c8>\u001b[0m in \u001b[0;36m<module>\u001b[0;34m\u001b[0m\n\u001b[0;32m----> 1\u001b[0;31m \u001b[0mdataset\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m",
      "\u001b[0;31mNameError\u001b[0m: name 'dataset' is not defined"
     ]
    }
   ],
   "source": [
    "dataset"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 17,
   "metadata": {},
   "outputs": [],
   "source": [
    "train = xr.open_dataset('../tcdata/enso_round1_train_20210201/CMIP_label.nc')\n",
    "train_nino= train['nino'].values"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 18,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "(4645, 36)"
      ]
     },
     "execution_count": 18,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "train_nino.shape"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 24,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "(55740, 24, 72)"
      ]
     },
     "execution_count": 24,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": 56,
   "metadata": {},
   "outputs": [],
   "source": [
    "train = xr.open_dataset('../tcdata/enso_round1_train_20210201/CMIP_train.nc')\n",
    "label = xr.open_dataset('../tcdata/enso_round1_train_20210201/CMIP_label.nc')    \n",
    "train_sst = train['sst'].values\n",
    "train_sst= np.concatenate((train_sst[:151*5],train_sst[151*9:151*12],train_sst[151*13:]))  \n",
    "train_sst=split_month(train_sst,1000)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 57,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "(11976, 12, 24, 72)"
      ]
     },
     "execution_count": 57,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "train_sst.shape"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 64,
   "metadata": {},
   "outputs": [],
   "source": [
    "a=split_month_label(train_label,100)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 70,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "(3890, 36)"
      ]
     },
     "execution_count": 70,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "train_label.shape"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 114,
   "metadata": {},
   "outputs": [],
   "source": [
    "size=100\n",
    "temp=train_label[:size,0:12]\n",
    "temp=temp.reshape(size*12)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 115,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "(1200,)"
      ]
     },
     "execution_count": 115,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "temp.shape"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 121,
   "metadata": {},
   "outputs": [],
   "source": [
    "aa=[temp[i+12:i+36] for i in range(size*12-40)]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 122,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "(1160, 24)"
      ]
     },
     "execution_count": 122,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "np.array(aa).shape"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": 117,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "1176"
      ]
     },
     "execution_count": 117,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "len(aa)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 120,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "(13,)"
      ]
     },
     "execution_count": 120,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": 63,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "(11976, 12, 24, 72)"
      ]
     },
     "execution_count": 63,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "split_month(train_sst,1000).shape"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 87,
   "metadata": {},
   "outputs": [],
   "source": [
    "def split_month(array,size):#input shape: :,36,24,72\n",
    "    temp=array[:size,0:12,:,:]\n",
    "    temp=temp.reshape(size*12,24,72)\n",
    "    temp2=np.array([temp[i:i+12,:,:] for i in range(size*12-24)])\n",
    "    return temp2\n",
    "\n",
    "def split_month_label(array,size):#input shape: :,24\n",
    "    temp=array[:size,0:12]\n",
    "    temp=temp.reshape(size*12)\n",
    "    temp2=np.array([temp[i+12:i+36] for i in range(size*12-24)])\n",
    "    return temp2\n",
    "    \n",
    "def load_data2():\n",
    "    # CMIP data    \n",
    "    size1=500\n",
    "    train = xr.open_dataset('../tcdata/enso_round1_train_20210201/CMIP_train.nc')\n",
    "    label = xr.open_dataset('../tcdata/enso_round1_train_20210201/CMIP_label.nc')    \n",
    "    train_sst = train['sst'].values\n",
    "    train_sst= np.concatenate((train_sst[:151*5],train_sst[151*9:151*12],train_sst[151*13:]))  \n",
    "    train_sst=split_month(train_sst,size1)\n",
    "    train_t300 = train['t300'].values\n",
    "    train_t300= np.concatenate((train_t300[:151*5],train_t300[151*9:151*12],train_t300[151*13:]))\n",
    "    train_t300=split_month(train_t300,size1)\n",
    "    train_ua = train['ua'].values\n",
    "    train_ua= np.concatenate((train_ua[:151*5],train_ua[151*9:151*12],train_ua[151*13:])) \n",
    "    train_ua=split_month(train_ua,size1)\n",
    "    train_va = train['va'].values\n",
    "    train_va= np.concatenate((train_va[:151*5],train_va[151*9:151*12],train_va[151*13:]))\n",
    "    train_va=split_month(train_va,size1)\n",
    "    train_label = label['nino'].values\n",
    "    train_label= np.concatenate((train_label[:151*5],train_label[151*9:151*12],train_label[151*13:]))\n",
    "    train_label=split_month_label(train_label,size1)\n",
    "    \n",
    "    #train_ua = np.nan_to_num(train_ua)#缺失值补0\n",
    "    #train_va = np.nan_to_num(train_va)\n",
    "    #train_t300 = np.nan_to_num(train_t300)\n",
    "    #train_sst = np.nan_to_num(train_sst)\n",
    "\n",
    "    # SODA data  \n",
    "    size2=100\n",
    "    train2 = xr.open_dataset('../tcdata/enso_round1_train_20210201/SODA_train.nc')\n",
    "    label2 = xr.open_dataset('../tcdata/enso_round1_train_20210201/SODA_label.nc')\n",
    "    \n",
    "    train_sst2 = train2['sst'].values  # (3890, 12, 24, 72)\n",
    "    train_sst2=split_month(train_sst2,size2)\n",
    "    train_t3002 = train2['t300'].values\n",
    "    train_t3002=split_month(train_t3002,size2)\n",
    "    train_ua2 = train2['ua'].values\n",
    "    train_ua2=split_month(train_ua2,size2)\n",
    "    train_va2 = train2['va'].values\n",
    "    train_va2=split_month(train_va2,size2)\n",
    "    train_label2 = label2['nino'].values\n",
    "    train_label2=split_month_label(train_label2,size2)\n",
    "\n",
    "    print('Train samples: {}, Valid samples: {}'.format(len(train_label), len(train_label2)))\n",
    "\n",
    "    dict_train = {\n",
    "        'sst':train_sst,\n",
    "        't300':train_t300,\n",
    "        'ua':train_ua,\n",
    "        'va': train_va,\n",
    "        'label': train_label}\n",
    "    dict_valid = {\n",
    "        'sst':train_sst2,\n",
    "        't300':train_t3002,\n",
    "        'ua':train_ua2,\n",
    "        'va': train_va2,\n",
    "        'label': train_label2}\n",
    "    train_dataset = EarthDataSet(dict_train)\n",
    "    valid_dataset = EarthDataSet(dict_valid)\n",
    "    return train_dataset, valid_dataset\n",
    "\n",
    "class EarthDataSet(Dataset):\n",
    "    def __init__(self, data):\n",
    "        self.data = data\n",
    "\n",
    "    def __len__(self):\n",
    "        return len(self.data['sst'])\n",
    "\n",
    "    def __getitem__(self, idx):   \n",
    "        return (self.data['sst'][idx], self.data['t300'][idx], self.data['ua'][idx], self.data['va'][idx]), self.data['label'][idx]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 66,
   "metadata": {},
   "outputs": [],
   "source": [
    "train = xr.open_dataset('../tcdata/enso_round1_train_20210201/CMIP_train.nc')\n",
    "label = xr.open_dataset('../tcdata/enso_round1_train_20210201/CMIP_label.nc')    \n",
    "train_sst = train['sst'].values\n",
    "train_sst= np.concatenate((train_sst[:151*5],train_sst[151*9:151*12],train_sst[151*13:])) "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 91,
   "metadata": {},
   "outputs": [],
   "source": [
    "a=train_sst[:10]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 92,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "(10, 36, 24, 72)"
      ]
     },
     "execution_count": 92,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "a.shape"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 83,
   "metadata": {},
   "outputs": [],
   "source": [
    "#a=a.reshape(1,10,36,24,72)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 85,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "(1, 10, 36, 24, 72)"
      ]
     },
     "execution_count": 85,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "#a.shape"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 96,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "(2, 10, 36, 24, 72)"
      ]
     },
     "execution_count": 96,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "b=np.array([a,a])\n",
    "b.shape"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 59,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Train samples: 5976, Valid samples: 1176\n"
     ]
    }
   ],
   "source": [
    "res1,res2=load_data2()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "pytorch",
   "language": "python",
   "name": "pytorch"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.7.6"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}
